#!/usr/bin/python3
# -*- coding:utf-8 -*-
import torch
import numpy as np
from einops import repeat

def gen_golden_data_simple():
    # Parse tiling data in ../main.cpp
    num_rows = 0
    num_columns = 0
    batch_shape = []
    with open("./main.cpp", "r") as f:
        lines = f.readlines()
        for line in lines:
            if "uint32_t blockDim" in line:
                block_dim = int(line.split()[-1].split(";")[0])
            if "uint32_t num_rows" in line:
                num_rows = int(line.split()[-1].split(";")[0])
            if "uint32_t num_columns" in line:
                num_columns = int(line.split()[-1].split(";")[0])
            if "std::vector<uint32_t> batch_shape" in line:
                batch_shape = list(map(int, line.replace("{", "").replace("}", "").split(" = ")[-1].split(";")[0].split(",")))

    with open("./common.h", "r") as f:
        lines = f.readlines()
        # Parse typedef float DTYPE_Y;
        for line in lines:
            if "typedef _Float32 DTYPE_Y" in line:
                dtype = np.float32
                break
            elif "typedef _Float16 DTYPE_Y" in line:
                dtype = np.float16
                break
            elif "typedef int32_t DTYPE_Y" in line:
                dtype = np.int32
                break
            elif "typedef double DTYPE_Y" in line:
                dtype = np.float64
                break
    
    print(f"block_dim: {block_dim}, num_rows: {num_rows}, num_columns: {num_columns}, batch_shape: {batch_shape}, dtype: {dtype}")

    input_x = torch.randn(block_dim, *batch_shape, num_rows, num_columns).numpy().astype(dtype)
    golden = torch.eye(
        n=num_rows,
        m=num_columns,
    )
    golden = repeat(golden, "n m -> b n m", b=torch.prod(torch.tensor(batch_shape)) * block_dim).numpy().astype(dtype)

    input_x.tofile("./input/input_x.bin")
    golden.tofile("./output/golden.bin")

if __name__ == "__main__":
    gen_golden_data_simple()
